# -*- coding:utf-8 -*-
# 将连续内存镜像获得的json文件转换为可以在blender中使用的json文件

"""
Generates network coordinates using a layered layout.

The input arguments require a starting node that is assigned to
the starting layer (layer 0). Nodes are arranged by their distance
from this starting node.

"""

from __future__ import print_function

from random import uniform
from math import sqrt
from itertools import combinations, repeat
import numpy as np
from operator import add
import sys
import json
from itertools import takewhile
import re
from json import encoder
encoder.FLOAT_REPR = lambda o: format(o, '.3f')

'''
输出json格式内容
'''
def compress(obj):
    """Outputs json without whitespace."""
    return json.dumps(obj, sort_keys=True, separators=(",", ":"),
                      cls=CustomEncoder)


def dumps(obj):
    """Outputs json with formatting edits + object handling."""
    return json.dumps(obj, indent=4, sort_keys=True, cls=CustomEncoder)


class CustomEncoder(json.JSONEncoder):

    def encode(self, obj):
        """Fired for every object."""
        s = super(CustomEncoder, self).encode(obj)
        # If uncompressed, postprocess for formatting
        if len(s.splitlines()) > 1:
            s = self.postprocess(s)
        return s

    def postprocess(self, json_string):
        """Displays each entry on its own line."""
        is_compressing, is_hash, compressed, spaces = False, False, [], 0
        for row in json_string.split("\n"):
            if is_compressing:
                if (row[:spaces + 5] == " " * (spaces + 4) +
                        ("\"" if is_hash else "{")):
                    compressed.append(row.rstrip())
                elif (len(row) > spaces and row[:spaces] == " " * spaces and
                        re.match("[\]\}],?", row[spaces:].rstrip())):
                    compressed.append(row.rstrip())
                    is_compressing = False
                else:
                    compressed[-1] += " " + row.strip()
            else:
                compressed.append(row.rstrip())
                if any(a in row for a in ["edges", "nodes"]):
                    # Fix to handle issues that arise with empty lists
                    if "[]" in row:
                        continue
                    spaces = sum(1 for _ in takewhile(str.isspace, row))
                    is_compressing, is_hash = True, "{" in row
        return "\n".join(compressed)


'''
生成结点和边
'''

def pol2cart(rho, phi):
    """Converts radial coordinates to cartesian."""
    return [rho * np.cos(phi), rho * np.sin(phi)]


def run_forcing(edges, nodes, iterations=1000, force_strength=5.0,
                dampening=0.05, max_velocity=2.0, max_distance=50,
                is_3d=False):
    """Runs a force-directed-layout algorithm on the input graph.

    iterations - Number of FDL iterations to run in coordinate generation
    force_strength - Strength of Coulomb and Hooke forces
                     (edit this to scale the distance between nodes)
    dampening - Multiplier to reduce force applied to nodes
    max_velocity - Maximum distance a node can move in one step
    max_distance - The maximum distance considered for interactions
    """

    # Convert to a data-storing object and initialize some values
    d = 3 if is_3d else 2
    nodes = {n: {"velocity": nodes[n]["location"],
                 "force": [0.0] * d} for n in nodes}

    # Repeat n times (is there a more Pythonic way to do this?)
    for _ in repeat(None, iterations):

        # Add in Coulomb-esque node-node repulsive forces
        for node1, node2 in combinations(nodes.values(), 2):
            _coulomb(node1, node2, force_strength, max_distance)

        # And Hooke-esque edge spring forces
        for edge in edges:
            _hooke(nodes[edge["source"]], nodes[edge["target"]],
                   force_strength, max_distance)

        # Move by resultant force
        for node in nodes.values():
            # Constrain the force to the bounds specified by input parameter
            force = [_constrain(dampening * f, -max_velocity, max_velocity)
                     for f in node["force"]]
            # Update velocities and reset force
            node["velocity"] = [v + dv
                                for v, dv in zip(node["velocity"], force)]
            node["force"] = [0] * d

    # Clean and return
    for node in nodes.values():
        del node["force"]
        node["location"] = node["velocity"]
        del node["velocity"]
        # Even if it's 2D, let's specify three dimensions
        if not is_3d:
            node["location"] += [0.0]
    return nodes


def layer(edges, available_nodes, start_node="1", depth=0, xycoor=[0.0, 0.0],
          separation=10.0, rdistance=20.0, is_3d=True):
    """Runs a layered-layout algorithm on the input graph.
    This a recursive function that will run on each child node.

    Parameters to mess with:
    separation - the separation between the layers in terms of the z axis
    rdistance - if there is more than one child node,
        they are arrange radially from the parent.
    If running the spacing algorithm, ultimately this will not matter too much.
    """
    # initialize the locations and convert the nodes to a dictionary of
    # dictionaries
    child_nodes = []
    nodes = {}

    # remove the current node from the list of available ones
    available_nodes.remove(start_node)

    # set the parent node higher
    nodes[start_node] = {}
    nodes[start_node]["depth"] = depth
    nodes[start_node]["location"] = xycoor + [-1 * separation * depth]

    # get the children of the parent
    for e in edges:
        if e["source"] == start_node:
            child_nodes.append(e["target"])
        if e["target"] == start_node:
            child_nodes.append(e["source"])

    # set the children equidistant from each other for initial positioning and
    # layer from them
    if len(child_nodes) > 1:
        angle_btw = [2 * np.pi / len(child_nodes) * i
                     for i in range(0, len(child_nodes))]
        for i in range(0, len(child_nodes)):
            if(child_nodes[i] in available_nodes):
                # center around the parent node
                location_vector = map(add, pol2cart(
                    rdistance, angle_btw[i]), xycoor)
                nodes.update(layer(edges, available_nodes, child_nodes[
                             i], depth + 1, location_vector))
    elif len(child_nodes) == 1:
        if(child_nodes[0] in available_nodes):
            nodes.update(layer(edges, available_nodes,
                               child_nodes[0], depth + 1, xycoor))

    return nodes


def space(input_nodes, nodes_of_interest, force_iter=10):
    """
    This function is designed to space the crowded layers.
    It is very hacked, and improved solutions would be welcomed.

    The general premise is to treat the nodes in a given layer as a subnetwork
     and make fake connections. Using the fake connections, the forcing code
     written by Pat Fuller is applied.

    Parameters -
    force_iter - changes the magnitude of the force applied to the particles
    """

    # get the locations of the nodes of interest in 2-D
    # makes a subdictionary for the particular layer
    nodes = dict((k, input_nodes[k])
                 for k in nodes_of_interest if k in input_nodes)

    # a matrix to store distances
    dist_mat = np.zeros(shape=(len(nodes.keys()), len(nodes.keys())))

    # store the distances in the matrix
    for node1, node2 in combinations(nodes.keys(), 2):
        delta = [x2 - x1 for x1,
                 x2 in zip(nodes[node1]["location"], nodes[node2]["location"])]
        distance = sqrt(sum(d ** 2 for d in delta))
        dist_mat[nodes_of_interest.index(
            node1), nodes_of_interest.index(node2)] = distance
        dist_mat[nodes_of_interest.index(
            node2), nodes_of_interest.index(node1)] = distance

    # find the two furthest nodes
    [max_loc, temp] = np.where(dist_mat == dist_mat.max())
    [forward_node, reverse_node] = max_loc

    # create edges starting from the terminal nodes
    temp_edges = []
    forward_mat = dist_mat.copy()
    reverse_mat = dist_mat.copy()

    # walk through the network and makes temporary connections
    for i in range(0, len(nodes.keys())):
        next_node = forward_mat[forward_node, :].argsort()[1]
        forward_mat[forward_node, :] = dist_mat.max() + 1
        forward_mat[:, forward_node] = dist_mat.max() + 1
        temp_edges.append({"source": nodes_of_interest[
                forward_node], "target": nodes_of_interest[next_node]})
        forward_node = next_node

        next_node = reverse_mat[reverse_node, :].argsort()[1]
        reverse_mat[reverse_node, :] = dist_mat.max() + 1
        reverse_mat[:, reverse_node] = dist_mat.max() + 1
        temp_edges.append({"source": nodes_of_interest[
                reverse_node], "target": nodes_of_interest[next_node]})
        reverse_node = next_node

    nodes = run_forcing(temp_edges, nodes, iterations=force_iter)

    # update the x,y values for the forced nodes
    for node in nodes.keys():
        input_nodes[node]["location"][0:2] = nodes[node]["location"][0:2]

    return input_nodes


def _coulomb(n1, n2, k, r):
    """Calculates Coulomb forces and updates node data."""
    # Get relevant positional data
    delta = [x2 - x1 for x1, x2 in zip(n1["velocity"], n2["velocity"])]
    distance = sqrt(sum(d ** 2 for d in delta))

    # If the deltas are too small, use random values to keep things moving
    if distance < 0.1:
        delta = [uniform(0.1, 0.2) for _ in repeat(None, 3)]
        distance = sqrt(sum(d ** 2 for d in delta))

    # If the distance isn't huge (ie. Coulomb is negligible), calculate
    if distance < r:
        force = (k / distance) ** 2
        n1["force"] = [f - force * d for f, d in zip(n1["force"], delta)]
        n2["force"] = [f + force * d for f, d in zip(n2["force"], delta)]


def _hooke(n1, n2, k, r):
    """Calculates Hooke spring forces and updates node data."""
    # Get relevant positional data
    delta = [x2 - x1 for x1, x2 in zip(n1["velocity"], n2["velocity"])]
    distance = sqrt(sum(d ** 2 for d in delta))

    # If the deltas are too small, use random values to keep things moving
    if distance < 0.1:
        delta = [uniform(0.1, 0.2) for _ in repeat(None, 3)]
        distance = sqrt(sum(d ** 2 for d in delta))

    # Truncate distance so as to not have crazy springiness
    distance = min(distance, r)

    # Calculate Hooke force and update nodes
    force = (distance ** 2 - k ** 2) / (distance * k)
    n1["force"] = [f + force * d for f, d in zip(n1["force"], delta)]
    n2["force"] = [f - force * d for f, d in zip(n2["force"], delta)]


def _constrain(value, min_value, max_value):
    """Constrains a value to the inputted range."""
    return max(min_value, min(value, max_value))


if __name__ == "__main__":
    # threshold for the number of nodes in a given layer for spacing algorithm
    # to be used
    threshold_for_spacing = 20

    try:
        with open(sys.argv[-1]) as in_file:
            edges = json.load(in_file)
    except IOError:
        edges = json.loads(sys.argv[-1])

    # get the starting node
    starting_node = "4:System"

    # Convert to internal representation
    # edges = [{"source": str(s), "target": str(t)} for s, t in edges]
    name_over60 = {}
    for link in edges["links"]:
        if link["target_group"] == 7:
            link["target"] = link["target"].replace("(", "")
            link["target"] = link["target"].replace(")", "")
            link["target"] = link["target"].replace("\'", "")
            sp = link["target"].split(", ")
            new_target = sp[0]
            fun = sp[1]
            link["target"] = new_target
            name_over60[new_target.split(":")[1]] = fun
    edges = edges["links"]

    # Handle additional args (Will need to change)
    kwargs = {"force_strength": 5.0, "is_3d": True}
    for i, arg in enumerate(sys.argv):
        if arg == "--force-strength":
            kwargs["force_strength"] = float(sys.argv[i + 1])

    # Handle additional args
    kwargs = {"separation": 5.0, "is_3d": True}

    # generate the available nodes
    available_nodes = set(e["source"]
                          for e in edges) | set(e["target"] for e in edges)

    # creates the initial arrangement for the nodes based on the starting node
    master_nodes = layer(edges, available_nodes,
                         start_node=starting_node, depth=0)

    # get the range of depth and counts for each layer
    depth_count = {}
    for node in master_nodes.keys():
        if master_nodes[node]["depth"] in depth_count.keys():
            depth_count[master_nodes[node]["depth"]] = depth_count[
                master_nodes[node]["depth"]] + 1
        else:
            depth_count[master_nodes[node]["depth"]] = 1

    # space out particular layers based on how many nodes are in a particular
    # layer
    for layer_of_interest, count in depth_count.items():
        if count > threshold_for_spacing:
            nodes_of_interest = []
            for node in master_nodes.keys():
                if master_nodes[node]["depth"] == layer_of_interest:
                    nodes_of_interest.append(node)

            master_nodes = space(
                master_nodes, nodes_of_interest, force_iter=10)

    # Convert to json
    network_str = dumps({"edges": edges, "nodes": master_nodes, "funOfName": name_over60})
    network = json.loads(network_str)

    # 添加add_color
    name_group = {}
    for edge in network['edges']:
        if edge['source'] not in name_group.keys():
            name_group[edge['source']] = edge['source_group']
        if edge['target'] not in name_group.keys():
            name_group[edge['target']] = edge['target_group']
    
    for key in network['nodes'].keys():
        if name_group[key] == 1:
            network['nodes'][key]['color'] = 'blue'
        elif name_group[key] == 2:
            network['nodes'][key]['color'] = 'yellow'
        elif name_group[key] == 3:
            network['nodes'][key]['color'] = 'green'
        elif name_group[key] == 4:
            network['nodes'][key]['color'] = 'indigo'
        elif name_group[key] == 5:
            network['nodes'][key]['color'] = 'orange'
        elif name_group[key] == 6:
            network['nodes'][key]['color'] = 'purple'
        elif name_group[key] == 7:
            network['nodes'][key]['color'] = 'gray'
    print(dumps(network))